import torch
import torch.nn as nn
from Testmodel import CNN
from dnn.datasets import CaptchaData
from torchvision.transforms import Compose, ToTensor, Resize


from project_setting import dnn_model_path

testpath = './datas'

source = [chr(i) for i in range(65, 65 + 26)]
alphabet = ''.join(source)

device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")

def predict(img_dir='./datas'):
    transforms = Compose([Resize((40, 120)), ToTensor()])
    dataset = CaptchaData(img_dir, transform=transforms)
    # dataset = MyDataSet()
    cnn = CNN()
    if torch.cuda.is_available():
        cnn = cnn.to(device)
        cnn.eval()
        cnn.load_state_dict(torch.load(dnn_model_path))
    else:
        cnn.eval()
        model = torch.load(dnn_model_path, map_location='cpu')
        cnn.load_state_dict(model)

    for k, (img, target) in enumerate(dataset):
        img = img.view(1, 3, 40, 120).to(device)
        output = cnn(img)

        output = output.view(-1, 26)
        target = target.view(-1, 26).to(device)
        output = nn.functional.softmax(output, dim=1)
        output = torch.argmax(output, dim=1)
        target = torch.argmax(target, dim=1)
        output = output.view(-1, 4)[0]
        a = ''.join([alphabet[i] for i in output.cpu().numpy()])

        break

    return a


if __name__ == "__main__":
    
    image = predict(img_dir=testpath)
    print(image)